#define _GNU_SOURCE      // 启用GNU扩展（clone、pivot_root等）
#include <errno.h>       // 错误码处理
#include <sched.h>       // clone系统调用
#include <stdio.h>       // 输入输出
#include <string.h>      // 字符串操作
#include <sys/wait.h>    // 等待子进程
#include <unistd.h>      // execvp、chdir等
#include <sys/mount.h>   // mount、umount、pivot_root
#include <sys/stat.h>    // mkdir
#include <fcntl.h>       // open

// 声明pivot_root函数以消除警告
extern int pivot_root(const char *new_root, const char *put_old);

#define STACKSIZE (1024 * 1024)  // 子进程栈大小（1MB，栈向下增长）
static char stack[STACKSIZE];    // 子进程栈空间

// 错误打印函数
void print_err(const char *reason) {
    fprintf(stderr, "Error %s: %s\n", reason, strerror(errno));
}

// 3. UID/GID映射配置（必须由父进程调用）
void setup_id_mapping(pid_t pid) {
    // 配置UID映射：容器内UID 0 → 主机UID（getuid()），映射范围1个ID
    char uid_map_path[256];
    snprintf(uid_map_path, sizeof(uid_map_path), "/proc/%d/uid_map", pid);
    int fd = open(uid_map_path, O_WRONLY);
    if (fd == -1) { print_err("open uid_map"); close(fd); return; }
    
    char uid_map[64];
    snprintf(uid_map, sizeof(uid_map), "0 %d 1", getuid());  // 格式：容器UID 主机UID 映射数量
    if (write(fd, uid_map, strlen(uid_map)) == -1)
        print_err("write uid_map");
    close(fd);

    // 配置GID映射（逻辑同UID）
    char gid_map_path[256];
    snprintf(gid_map_path, sizeof(gid_map_path), "/proc/%d/gid_map", pid);
    fd = open(gid_map_path, O_WRONLY);
    if (fd == -1) { print_err("open gid_map"); close(fd); return; }
    
    char gid_map[64];
    snprintf(gid_map, sizeof(gid_map), "0 %d 1", getgid());
    if (write(fd, gid_map, strlen(gid_map)) == -1)
        print_err("write gid_map");
    close(fd);
}

// 子进程执行逻辑（包含根目录切换、proc挂载、主机名设置）
int exec(void *args) {
    char **argv = args;
    const char *new_root = argv[0];  // 第1个参数：新根目录路径
    char **cmd = &argv[1];           // 第2个及以后参数：要执行的命令（如/bin/sh）

    // 1. 使用chroot切换根目录
    // 步骤1：先切换到新根目录的父目录，确保后续能访问新根
    char parent_dir[256];
    strncpy(parent_dir, new_root, sizeof(parent_dir)-1);
    char *last_slash = strrchr(parent_dir, '/');
    if (last_slash && last_slash != parent_dir) {
        *last_slash = '\0';  // 截断到父目录
    } else {
        // 处理根目录情况
        strncpy(parent_dir, "/", sizeof(parent_dir)-1);
    }

    // 步骤2：切换到父目录
    if (chdir(parent_dir) != 0) {
        print_err("chdir to parent directory");
        return 1;
    }

    // 步骤3：执行chroot切换根目录
    if (chroot(new_root) != 0) {
        print_err("chroot");
        return 1;
    }

    // 步骤4：切换到新的根目录
    if (chdir("/") != 0) {
        print_err("chdir to new /");
        return 1;
    }

    // 2. 重新挂载/proc（确保显示新PID命名空间内的进程）
    if (mount("proc", "/proc", "proc", 0, "") != 0) {
        print_err("mount /proc");
        return 1;
    }

    // 4. 设置新主机名（UTS命名空间隔离生效）
    const char *hostname = "container-host";
    if (sethostname(hostname, strlen(hostname)) != 0) {
        print_err("sethostname");
        return 1;
    }

    // 执行目标命令（替换当前进程）
    if (execvp(cmd[0], cmd) != 0) {
        print_err("exec command");
        return 1;
    }
    return 0;
}

int main(int argc, char **argv) {
    // 参数检查：需传入「新根目录路径」和「要执行的命令」（如 ./container ./newroot /bin/sh）
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <new-root-path> <command> [args...]\n", argv[0]);
        fprintf(stderr, "Example: %s ./newroot /bin/sh\n", argv[0]);
        return 1;
    }

    // 命名空间标志：启用所有必要隔离（网络、UTS、挂载、IPC、PID、用户、CGROUP）
    const int flags = CLONE_NEWNET | CLONE_NEWUTS | CLONE_NEWNS |
                      CLONE_NEWIPC | CLONE_NEWPID | CLONE_NEWUSER |
                      CLONE_NEWCGROUP | SIGCHLD;  // SIGCHLD：子进程退出时通知父进程

    // 创建带命名空间的子进程（栈指针指向栈顶，因栈向下增长）
    pid_t pid = clone(exec, stack + STACKSIZE, flags, &argv[1]);
    if (pid < 0) {
        print_err("clone (create child)");
        return 1;
    }

    // 父进程：配置UID/GID映射（必须在子进程执行exec前调用）
    setup_id_mapping(pid);

    // 等待子进程结束，获取退出状态
    int status = 0;
    if (waitpid(pid, &status, 0) == -1) {
        print_err("waitpid");
        return 1;
    }

    return WEXITSTATUS(status);
}